(*
 *  MTASC - MotionTwin ActionScript2 Compiler
 *  Copyright (c)2004 Nicolas Cannasse
 *
 *  This program is free software; you can redistribute it and/or modify
 *  it under the terms of the GNU General Public License as published by
 *  the Free Software Foundation; either version 2 of the License, or
 *  (at your option) any later version.
 *
 *  This program is distributed in the hope that it will be useful,
 *  but WITHOUT ANY WARRANTY; without even the implied warranty of
 *  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 *  GNU General Public License for more details.
 *
 *  You should have received a copy of the GNU General Public License
 *  along with this program; if not, write to the Free Software
 *  Foundation, Inc., 59 Temple Place, Suite 330, Boston, MA  02111-1307  USA
 *)

open Expr
open Swf
open ExtHashtbl
open ExtString
open ExtList
open GenSwfJump
open GenSwfTools
open GenSwfGeneratorTools

let generate_local_var ftrace ctx (vname,_,vinit) =
	if used_in_block false vname ctx.cur_block || ctx.reg_count >= 250 then begin
		push ctx [VStr vname];
		Hashtbl.add ctx.locals vname { reg = 0; sp = ctx.stack };
		match vinit with
		| None -> write ctx ALocalVar
		| Some v ->
			generate_val ftrace ctx v;
			write ctx ALocalAssign
	end else begin
		ctx.reg_count <- ctx.reg_count + 1;
		let r = ctx.reg_count in
		Hashtbl.add ctx.locals vname { reg = r; sp = ctx.stack };
		match vinit with
		| None -> 
			()
		| Some v ->
			generate_val ftrace ctx v;
			setvar ctx (VarReg r)
	end

let gen_forins ctx all =
	for i = 1 to (if all then ctx.forins else 1) do
		push ctx [VNull];
		write ctx AEqual;
		write ctx ANot;
		write ctx (ACondJump (-4));
	done

let rec generate_expr ftrace ctx (e,p) =
	match e with
	| EFunction _ ->
		assert false
	| EVars (_,_,vl) ->
		List.iter (generate_local_var ftrace ctx) vl
	| EBlock el ->
		let block_end = open_block ctx (e,p) in
		List.iter (generate_expr ftrace ctx) el;
		block_end()
	| EFor (inits,conds,incrs,e) ->
		let block_end = open_block ctx e in
		List.iter (generate_expr ftrace ctx) inits;
		let test = jmp ctx in
		let start_pos = ctx.code_pos in
		let old_continue = ctx.continue_pos in
		let old_breaks = ctx.breaks in
		ctx.breaks <- [];
		ctx.continue_pos <- start_pos;
		ctx.opt_push <- false;
		List.iter (fun v -> generate_expr ftrace ctx (EVal v,null_pos)) incrs;
		test();
		let jumps = ref [] in
		List.iter (fun cond -> 
			generate_val ftrace ctx cond;
			write ctx ANot;
			jumps := cjmp ctx :: !jumps;
		) conds;
		generate_expr ftrace ctx e;
		do_jmp ctx start_pos;
		List.iter (fun j -> j()) !jumps;
		generate_breaks ctx old_breaks;
		ctx.continue_pos <- old_continue;
		block_end()
	| EForIn (decl,v,e) ->
		let block_end = open_block ctx e in
		generate_val ftrace ctx v;
		write ctx AEnum2;
		let start_pos = ctx.code_pos in
		let old_continue = ctx.continue_pos in
		let old_breaks = ctx.breaks in
		ctx.breaks <- [];
		ctx.continue_pos <- start_pos;
		ctx.opt_push <- false;
		ctx.forins <- ctx.forins + 1;
		write ctx (ASetReg 0);
		push ctx [VNull];
		write ctx AEqual;
		let jump_end = cjmp ctx in
		(match fst decl with
		| EVal ((EConst (Ident _),_) as x) ->
			let k = generate_access ftrace ctx x in
			push ctx [VReg 0];
			setvar ctx k
		| EVars (_,_,[(x,_,None)]) ->
			push ctx [VStr x];
			Hashtbl.add ctx.locals x { reg = 0; sp = ctx.stack };
			push ctx [VReg 0];
			write ctx ALocalAssign
		| _ ->
			error (pos decl));
		generate_expr ftrace ctx e;
		do_jmp ctx start_pos;
		let has_breaks = (ctx.breaks <> []) in
		generate_breaks ctx old_breaks;
		if has_breaks then gen_forins ctx false;
		jump_end();
		ctx.forins <- ctx.forins - 1;
		ctx.continue_pos <- old_continue;
		block_end()
	| EIf (v,e,eelse) ->
		generate_val ftrace ctx v;
		write ctx ANot;
		let jump_else = cjmp ctx in
		generate_expr ftrace ctx e;
		(match eelse with
		| None -> jump_else()
		| Some e ->
			let jump_end = jmp ctx in
			jump_else();
			generate_expr ftrace ctx e;
			jump_end())
	| EVal v ->
		let s = ctx.stack_size in
		generate_val ftrace ~retval:false ctx v;
		pop ctx (ctx.stack_size - s)
	| EWhile (v,e,flag) ->
		let jump_begin = (match flag with NormalWhile -> (fun()->()) | DoWhile -> jmp ctx) in
		let start_pos = ctx.code_pos in
		let old_continue = ctx.continue_pos in
		let old_breaks = ctx.breaks in
		ctx.breaks <- [];
		ctx.opt_push <- false;
		ctx.continue_pos <- start_pos;
		generate_val ftrace ctx v;
		write ctx ANot;
		let jump_end = cjmp ctx in
		jump_begin();
		generate_expr ftrace ctx e;
		do_jmp ctx start_pos;
		generate_breaks ctx old_breaks;
		ctx.continue_pos <- old_continue;
		jump_end()
	| EWith (v,e) ->
		generate_val ftrace ctx v;
		write ctx (AWith 0);
		let start = ctx.code_pos in
		generate_expr ftrace ctx e;
		let delta = ctx.code_pos - start in
		DynArray.set ctx.ops (start - 1) (AWith delta);
	| EBreak ->
		ctx.breaks <- jmp ctx :: ctx.breaks
	| EContinue ->
		do_jmp ctx ctx.continue_pos
	| EReturn None ->
		gen_forins ctx true;
		write ctx (APush [PUndefined]);
		write ctx AReturn	
	| EReturn (Some v) ->
		gen_forins ctx true;
		generate_val ftrace ctx v;
		write ctx AReturn
	| ESwitch (v,cases) ->
		generate_val ftrace ctx v;
		write ctx (ASetReg 0);
		let old_breaks = ctx.breaks in
		let first_case = ref true in
		let def_pos = ref (fun () -> ()) in
		ctx.breaks <- [];
		let cases = List.map (fun (v,e) ->
			match v with
			| None ->
				(fun () -> (!def_pos)(); def_pos := (fun() -> ())) , e
			| Some v -> 
				if !first_case then
					first_case := false
				else
					push ctx [VReg 0];
				generate_val ftrace ctx v;
				write ctx APhysEqual;
				cjmp ctx , e
		) cases in
		def_pos := jmp ctx;
		List.iter (fun (j,e) ->
			j();
			generate_expr ftrace ctx e
		) cases;
		(!def_pos)();
		generate_breaks ctx old_breaks
	| ETry (e,cl,fo) ->
		let tdata = {
			tr_style = TryRegister 0;
			tr_trylen = 0;
			tr_catchlen = None;
			tr_finallylen = None;
		} in
		write ctx (ATry tdata);
		let p = ctx.code_pos in
		generate_expr ftrace ctx e;
		let jump_end = jmp ctx in
		tdata.tr_trylen <- ctx.code_pos - p;
		let p = ctx.code_pos in
		let end_throw = ref true in
		let first_catch = ref true in
		let jumps = List.map (fun (name,t,e) ->
			Hashtbl.add ctx.locals name { reg = 0; sp = ctx.stack };
			let next_catch = (match t with
			| None -> 
				end_throw := false;
				write ctx APop;
				push ctx [VStr name;VReg 0];
				write ctx ALocalAssign;
				generate_expr ftrace ctx e;
				(fun() -> ())
			| Some t ->
				if not !first_catch then write ctx APop;
				getvar ctx (generate_access ftrace ctx (EStatic t,pos e));
				push ctx [VReg 0];
				write ctx ACast;
				write ctx ADup;
				push ctx [VNull];
				write ctx AEqual;
				let c = cjmp ctx in
				push ctx [VStr name];
				write ctx ASwap;
				write ctx ALocalAssign;
				generate_expr ftrace ctx e;
				c
			) in
			first_catch := false;
			let j = jmp ctx in
			next_catch();
			Hashtbl.remove ctx.locals name;
			j
		) !cl in
		if !end_throw && !cl <> [] then begin
			write ctx APop;
			push ctx [VReg 0];
			write ctx AThrow;
		end;
		if !cl <> [] then tdata.tr_catchlen <- Some (ctx.code_pos - p);
		List.iter (fun j -> j()) jumps;
		jump_end();
		(match fo with
		| None -> ()
		| Some e ->
			let p = ctx.code_pos in
			generate_expr ftrace ctx e;
			tdata.tr_finallylen <- Some (ctx.code_pos - p))

let super_call = EVal (ECall ((EConst (Ident "super"),null_pos),[]),null_pos) , null_pos

let generate_function ftrace ?(constructor=false) ctx f =
	match f.fexpr with
	| None -> ()
	| Some fexpr ->
		let old_name = ctx.curmethod in
		let stack_base , old_nregs = ctx.stack , ctx.reg_count in
		let have_super = used_in_block true "super" fexpr in
		let reg_super = have_super || (constructor && Class.superclass ctx.current <> None) in
		let old_forin = ctx.forins in
		ctx.reg_count <- (if reg_super then 2 else 1);
		if f.fname <> "" then ctx.curmethod <- f.fname;
		ctx.forins <- 0;
		ctx.stack <- ctx.stack + 1;
		let args = List.map (fun (aname,_) ->
			let r = 
				(if used_in_block false aname fexpr then
					0
				else begin
					ctx.reg_count <- ctx.reg_count + 1;
					ctx.reg_count
				end)
			in
			Hashtbl.add ctx.locals aname { reg = r; sp = ctx.stack };
			r , aname
		) f.fargs in
		let fdone = func ctx args reg_super (used_in_block true "arguments" fexpr) in
		if constructor && Class.superclass ctx.current <> None && not have_super then generate_expr ftrace ctx super_call;
		generate_expr ftrace ctx fexpr;
		if f.fgetter = Setter then begin
			push ctx [VInt 0;VThis;VStr ("__get__"^f.fname)];
			call ctx VarObj 0;
			write ctx AReturn;
		end;
		fdone (ctx.reg_count + 1);
		clean_stack ctx stack_base;
		ctx.forins <- old_forin;
		ctx.reg_count <- old_nregs;
		ctx.curmethod <- old_name

let generate_super_bindings ftrace ctx =
	Hashtbl.iter (fun (path,fname) flag ->
		if not flag then begin
			Hashtbl.replace ctx.super_bindings (path,fname) true;
			let ident = super_binding_ident path fname in
			push ctx [VStr ident];
			let k = generate_access ftrace ctx (EStatic path,null_pos) in
			getvar ctx k;
			push ctx [VStr "prototype"];
			getvar ctx VarObj;
			push ctx [VStr fname];
			getvar ctx VarObj;
			setvar ctx VarStr;
		end;
	) ctx.super_bindings

let generate_class_code ftrace enable_main ctx clctx h =
	let cpath , cname = Class.path clctx in
	getvar ctx (generate_access ftrace ctx (EStatic (cpath,cname),null_pos));
	write ctx ANot;
	write ctx ANot;
	let jump_end_def = cjmp ctx in
	if not (Hashtbl.mem h cpath) then begin
		generate_package_register ctx cpath;
		Hashtbl.add h cpath ();
	end;
	let k = generate_package ctx cpath in
	push ctx [VStr cname];
	(match Class.constructor clctx with
	| None -> 
		let fdone = func ctx [] true false in
		(match Class.superclass clctx with
		| None -> ()
		| Some _ -> generate_expr ftrace ctx super_call);
		fdone 3
	| Some f ->
		generate_function ftrace ~constructor:true ctx f);
	write ctx (ASetReg 0);
	setvar ctx k;
	(match Class.superclass clctx with
	| None -> ()
	| Some csuper when !version = Some 6 ->	
		(* myclass.prototype.__proto__ = superclass.prototype *)
		push ctx [VReg 0; VStr "prototype"];
		getvar ctx VarObj;
		push ctx [VStr "__proto__"];
		getvar ctx (generate_access ftrace ctx (EStatic (Class.path csuper),null_pos));
		push ctx [VStr "prototype"];
		getvar ctx VarObj;
		setvar ctx VarObj;
		(* myclass.prototype.__constructor__ = superclass *)
		push ctx [VReg 0; VStr "prototype"];
		getvar ctx VarObj;
		push ctx [VStr "__constructor__"];
		getvar ctx (generate_access ftrace ctx (EStatic (Class.path csuper),null_pos));
		setvar ctx VarObj
	| Some csuper ->
		push ctx [VReg 0];
		getvar ctx (generate_access ftrace ctx (EStatic (Class.path csuper),null_pos));
		write ctx AExtends);
	push ctx [VReg 0; VStr "prototype"];
	getvar ctx VarObj;
	write ctx (ASetReg 1);
	write ctx APop;
	let getters = Hashtbl.create 0 in
	List.iter (fun f ->
		match f.fexpr with
		| None -> ()
		| Some _ ->
			push ctx [VReg (if f.fstatic = IsMember then 1 else 0)];
			let name = (match f.fgetter with
				| Normal ->  
					if f.fname = "main" && f.fstatic = IsStatic && !enable_main then begin
						match !(ctx.main) with
						| None -> ctx.main := Some (Class.path clctx);
						| Some path -> failwith ("Duplicate main entry point : " ^ s_type_path path ^ " and " ^ s_type_path (Class.path clctx))
					end;					
					
					f.fname
				| Getter -> 
					Hashtbl.add getters (f.fname,Getter,f.fstatic) ();
					"__get__" ^ f.fname
				| Setter -> 
					Hashtbl.add getters (f.fname,Setter,f.fstatic) ();
					"__set__" ^ f.fname)
			in
			push ctx [VStr name];
			generate_function ftrace ctx f;
			setvar ctx VarObj;
	) (Class.methods clctx);
	let dones = Hashtbl.create 0 in
	Hashtbl.iter (fun (name,get,stat) _ ->
		if Hashtbl.mem dones (name,get,stat) then
			()
		else
		let reg = (if stat = IsMember then 1 else 0) in
		let getter = (get = Getter || Hashtbl.mem getters (name,Getter,stat)) in
		let setter = (get = Setter || Hashtbl.mem getters (name,Setter,stat)) in
		let no_getset = AFunction { f_name = ""; f_args = []; f_codelen = 0 } in

		Hashtbl.add dones (name,Getter,stat) ();
		Hashtbl.add dones (name,Setter,stat) ();
		if setter then begin
			push ctx [VReg reg; VStr ("__set__" ^ name)];
			getvar ctx VarObj;
		end else
			write ctx no_getset;			
		if getter then begin
			push ctx [VReg reg; VStr ("__get__" ^ name)];
			getvar ctx VarObj;
		end else
			write ctx no_getset;
		push ctx [VStr name; VInt 3];
		push ctx [VReg reg; VStr "addProperty"];
		call ctx VarObj 3;
		write ctx APop;
	) getters;
	List.iter (fun cintf ->
		getvar ctx (generate_access ftrace ctx (EStatic (Class.path cintf),null_pos));
	) (Class.interfaces clctx);
	let nintf = List.length (Class.interfaces clctx) in
	if nintf > 0 then begin
		push ctx [VInt nintf; VReg 0];
		write ctx AImplements;
		ctx.stack_size <- ctx.stack_size - nintf;
	end;
	push ctx [VInt 1; VNull; VReg 1; VInt 3; VStr "ASSetPropFlags"];
	call ctx VarStr 3;
	write ctx APop;
	List.iter (fun (name,stat,v) ->
		push ctx [VReg (if stat = IsMember then 1 else 0); VStr name];
		generate_val ftrace ctx v;
		setvar ctx VarObj;
	) (Class.initvars clctx);
	generate_super_bindings ftrace ctx;
	jump_end_def()	
